# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo
import time

import dateutil.parser
from recordclass import recordclass

from coin.exchange.okex.kr_rest.product import OkexProduct, OkexCurrency
from coin.exchange.base.kr_rest.private_client_base import (PrivateClientBase,
                                                            ExWalletPrivateClientBase)

import coin.proto.coin_order_enums_pb2 as coin_order

from coin.exchange.okex_v3.kr_rest.native_private_client import (
    OkexNativePrivateClient,
    OkexOrderStatus,
)
from coin.proto.coin_query_pb2 import (
    AccountBalance,
    CurrencyBalance,
    CurrencyTransferBundle,
    CurrencyTransferElement,
    ProductOrderBundle,
    ProductOrderElement,
    ProductFillBundle,
    ProductFillElement,
)
from coin.exchange.okex_futures_v3.kr_rest.private_client \
    import convert_okex_exec_type_to_proto_fill_type

TradeFeeDetail = recordclass('TradeFeeDetail',
                             ['fee_currency', 'fee', 'price', 'buy_qty', 'sell_qty'])


class OkexPrivateParser(object):
  ProductType = OkexProduct

  @staticmethod
  def parse_balance(json_msg):
    """
    spot account balance
    [
     {
        "frozen": "0",
        "hold": "0",
        "id": "8357757",
        "currency": "USDT",
        "balance": "30.694669486635059",
        "available": "30.694669486635059",
        "holds": "0"
      },
      ....
    ]
    wallet balance
    [
      {
        "balance": 0.0,
        "available": 0.0,
        "currency": "QTUM",
        "hold": 0.0
      },
      ....
    ]
    """
    currency_balance_list = []
    for balance in json_msg:
      currency = OkexCurrency.FromStrNativeCurrencyNothrow(balance['currency'])
      if currency is None:
        continue

      if 'frozen' in balance:
        assert balance['frozen'] == balance['hold'], balance
      currency_balance = CurrencyBalance(
          currency=currency.currency,
          currency_native=currency.native_currency,
          available=float(balance['available']),
          total=float(balance['balance']),
          hold=float(balance['hold']),
      )
      currency_balance_list.append(currency_balance)

    return AccountBalance(
        exchange="Okex",
        each_balance=currency_balance_list,
    )

  @staticmethod
  def parse_margin_balance(json_msg):
    currency_balances = {}
    for balance in json_msg:
      for key, value in balance.items():
        if not key.startswith('currency'):
          continue
        native_currency = key.split(':')[-1]
        currency = OkexCurrency.FromStrNativeCurrencyNothrow(native_currency)
        if currency is None:
          continue
        if native_currency not in currency_balances:
          currency_balances[native_currency] = CurrencyBalance(
            currency=currency.currency, currency_native=currency.native_currency)
        currency_balance = currency_balances[native_currency]
        currency_balance.available += float(value['available'])
        currency_balance.hold += float(value['hold'])
        currency_balance.total += float(value['balance']) - float(value['borrowed']) - float(value['lending_fee'])
        currency_balances[native_currency] = currency_balance
    return AccountBalance(
        exchange="Okex",
        each_balance=currency_balances.values(),
    )

  @staticmethod
  def parse_orders(json_msg, product):
    """
    [
    {
      "created_at": "2018-10-24T09:59:15Z",
      "executed_value": "0",
      "filled_notional": "0",
      "filled_size": "0",
      "funds": "",
      "instrument_id": "BTC-USDT",
      "notional": "",
      "order_id": "1680287676308480",
      "price": "5800",
      "product_id": "BTC-USDT",
      "side": "buy",
      "size": "0.002",
      "status": "open",
      "timestamp": "2018-10-24T09:59:15Z",
      "type": "limit"
    },
    ]
    """
    currency_order_list = []
    for info in json_msg:
      prod = OkexProduct.FromStrNativeProduct(info['instrument_id'])
      if prod != product:
        continue

      if info['side'] == 'sell':
        side = coin_order.SELL_ORDER
      elif info['side'] == 'buy':
        side = coin_order.BUY_ORDER
      else:
        raise ValueError(info['side'])

      timestamp = int(dateutil.parser.parse(info['timestamp']).timestamp() * 1e9)
      currency_order = ProductOrderElement(
          symbol=product.symbol,
          side=side,
          price=float(info['price']),
          qty=float(info['size']),
          order_id=str(info['order_id']),
          created_timestamp=timestamp,
      )
      currency_order_list.append(currency_order)
    return ProductOrderBundle(exchange='Okex',
                              symbol=product.symbol,
                              each_order=currency_order_list)

  @staticmethod
  def parse_fills(order_list_msg, fill_msgs, product):
    """
    order list is the same format as orders, but always status=filled.
    fill update format:
    [
    {
      "exec_type": "T",
      "fee": "0.018",
      "instrument_id": "ETH-USDT",
      "ledger_id": "1706",
      "order_id": "1798782957193216",
      "price": "90",
      "side": "points_fee",
      "size": "0",
      "timestamp": "2018-11-14T08:14.09000Z"
    }
    ]
    """
    fill_list = []
    for index, info in enumerate(order_list_msg):
      prod = OkexProduct.FromStrNativeProduct(info['instrument_id'])
      if prod != product:
        continue

      assert info['status'] == OkexOrderStatus.FILLED.value
      if info['side'] == 'sell':
        side = coin_order.SELL_ORDER
        fee_currency = product.quote.symbol
        fee_currency_reverse = product.base.symbol
      elif info['side'] == 'buy':
        side = coin_order.BUY_ORDER
        fee_currency = product.base.symbol
        fee_currency_reverse = product.quote.symbol
      else:
        raise ValueError(info['side'])

      fill_msg = fill_msgs[index]
      trade_fee_details = {
          coin_order.TAKER_FILL_TYPE: TradeFeeDetail(fee_currency, 0., 0., 0., 0.),
          coin_order.MAKER_FILL_TYPE: TradeFeeDetail(fee_currency, 0., 0., 0., 0.),
      }

      for item in fill_msg:
        prod = OkexProduct.FromStrNativeProduct(item['instrument_id'])
        assert prod == product
        fill_type = convert_okex_exec_type_to_proto_fill_type(item['exec_type'])
        if (side == coin_order.SELL_ORDER and item['side'] == info['side'] and float(item['fee']) > 0) or \
              (side == coin_order.BUY_ORDER and item['side'] != info['side'] and float(item['fee']) > 0):
          trade_fee_details[fill_type].fee_currency = fee_currency_reverse
        trade_fee_details[fill_type].fee -= float(item['fee'])
        trade_fee_details[fill_type].price = float(item['price'])
        qty = float(item['size'])
        if item['side'] == 'points_fee':
          fee_currency = 'OKDK'
        elif item['side'] == 'buy':
          trade_fee_details[fill_type].buy_qty += qty
        elif item['side'] == 'sell':
          trade_fee_details[fill_type].sell_qty += qty
        else:
          raise ValueError(item)

      for fill_type, trade_fee_detail in trade_fee_details.items():
        if trade_fee_detail.price == 0:
          continue
        if side == coin_order.BUY_ORDER:
          qty = trade_fee_detail.buy_qty
        elif side == coin_order.SELL_ORDER:
          qty = trade_fee_detail.sell_qty
        price = trade_fee_detail.price
        fee = trade_fee_detail.fee
        fee_currency = trade_fee_detail.fee_currency
        fill = ProductFillElement(symbol=product.symbol,
                                  side=side,
                                  price=price,
                                  qty=qty,
                                  order_id=str(info['order_id']),
                                  fee=fee,
                                  fee_currency=fee_currency,
                                  fill_type=fill_type)
        fill_list.append(fill)
    return ProductFillBundle(exchange='Okex', symbol=product.symbol, each_fill=fill_list)

  @staticmethod
  def parse_transfers(withdraw_update_msg, deposite_update_msg, currency):
    """
    Status of withdrawal. -3:pending cancel; -2: cancelled; -1: failed; 0:pending;
    1:sending; 2:sent; 3:awaiting email verification; 4:awaiting manual verification;
    5:awaiting identity verification
    """
    withdraw_status_map = {
        '-3': 'CANCEL_CONFIRMED',
        '-2': 'CANCEL_CONFIRMED',
        '-1': 'TRANSFER_REJECTED',
        '0': 'TRANSFER_ACCEPTED',
        '1': 'TRANSFER_ACCEPTED',
        '2': 'TRANSFER_CONFIRMED',
        '3': 'TRANSFER_ACCEPTED',
        '4': 'TRANSFER_ACCEPTED',
        '5': 'TRANSFER_ACCEPTED'
    }
    """
    Status of deposit 
    (0: waiting for confirmation; 1: deposit credited; 2: deposit successful)
    """
    deposite_status_map = {
        '0': 'TRANSFER_ACCEPTED', '1': 'TRANSFER_ACCEPTED', '2': 'TRANSFER_CONFIRMED'
    }
    transfer_list = []
    for transfer_elem in withdraw_update_msg:
      assert currency.native_currency in transfer_elem['fee'], transfer_elem['fee']
      transfer = CurrencyTransferElement(
          currency=currency.currency,
          type='WITHDRAW',
          from_address=transfer_elem['from'],
          to_address=transfer_elem['to'],
          qty=float(transfer_elem['amount']),
          fee=float(transfer_elem['fee'].replace(currency.native_currency, '')),
          external_transfer_id=transfer_elem['txid'],
          status=withdraw_status_map[transfer_elem['status']],
          create_timestamp=int(dateutil.parser.parse(transfer_elem['timestamp']).timestamp() * 1e9))
      transfer_list.append(transfer)
    for transfer_elem in deposite_update_msg:
      transfer = CurrencyTransferElement(
          currency=currency.currency,
          type='DEPOSIT',
          from_address=transfer_elem['from'],
          to_address=transfer_elem['to'],
          qty=float(transfer_elem['amount']),
          external_transfer_id=transfer_elem['txid'],
          status=deposite_status_map[transfer_elem['status']],
          create_timestamp=int(dateutil.parser.parse(transfer_elem['timestamp']).timestamp() * 1e9))
      transfer_list.append(transfer)
    return CurrencyTransferBundle(exchange='Okex',
                                  market_type='Spot',
                                  currency=currency.currency,
                                  native_currency=currency.native_currency,
                                  each_transfer=transfer_list)


# noinspection PyAbstractClass
class OkexPrivateClient(PrivateClientBase):
  ProductType = OkexProduct

  def __init__(self, key_file):
    super().__init__()
    self.nprivc = OkexNativePrivateClient(key_file=key_file)

  def query_account_balance_impl(self):
    update = self.nprivc.query_accounts()
    update.msg = OkexPrivateParser.parse_balance(update.msg)
    return update

  def query_list_orders_impl(self, product):
    update = self.nprivc.query_order_list(status=[OkexOrderStatus.OPEN],
                                          instrument_id=product.native_symbol)
    update.msg = OkexPrivateParser.parse_orders(update.msg, product)
    return update

  def query_cancel_impl(self, product, order_id):
    raise NotImplementedError()

  def query_fills_impl(self, product):
    update = self.nprivc.query_order_list(status=[OkexOrderStatus.FILLED],
                                          instrument_id=product.native_symbol,
                                          limit=15)
    fill_updates = []
    for order in update.msg:
      prod = OkexProduct.FromStrNativeProduct(order['instrument_id'])
      if prod != product:
        continue
      order_id = str(order['order_id'])
      fill_update = self.nprivc.query_transaction_details(order_id=order_id,
                                                          instrument_id=product.native_symbol)
      time.sleep(0.2)  # rate limit 10 times per 2 seconds
      fill_updates.append(fill_update.msg)
    update.msg = OkexPrivateParser.parse_fills(update.msg, fill_updates, product)
    return update

  def query_account_transfer_impl(self, currency):
    currency = OkexCurrency.FromStr(currency)
    withdraw_update = self.nprivc.query_withdraw_history(currency.native_currency)
    deposit_update = self.nprivc.query_deposit_history(currency.native_currency)
    update = deposit_update
    update.msg = OkexPrivateParser.parse_transfers(withdraw_update.msg,
                                                   deposit_update.msg,
                                                   currency)
    return update

  def query_cancel_product_impl(self, product):
    raise NotImplementedError()

  def query_submit_impl(self, product, side, price, qty):
    raise NotImplementedError()


class OkexExWalletPrivateClient(ExWalletPrivateClientBase):
  ProductType = OkexProduct

  def __init__(self, key_file):
    super().__init__()
    self.nprivc = OkexNativePrivateClient(key_file=key_file)

  def query_account_balance_impl(self):
    update = self.nprivc.query_wallet()
    update.msg = OkexPrivateParser.parse_balance(update.msg)
    return update

class OkexMarginPrivateClient(OkexPrivateClient):

  def query_account_balance_impl(self):
    update = self.nprivc.query_margin_accounts()
    update.msg = OkexPrivateParser.parse_margin_balance(update.msg)
    return update
